from src.Dynamics import dynamics
from src.Kinematics import kinematics
from src.Plot_trajectory import plot_trajectory
from src.VelocityController import velocityController
import numpy as np
import math

def main():
    # 初始化动力学和运动学模型
    dyn = dynamics()
    kin = kinematics()
    # 设置无人机的质量和惯性矩阵
    mass = 1.0  # kg
    inertia = np.diag([0.1, 0.1, 0.2])  # kg*m^2, 假设的惯性矩阵
    dt=0.01
    
    # 初始化速度控制器
    v_controller = velocityController()


    points = np.empty((0, 3))  # 初始化一个空的二维数组，每行3列
    for _ in range(1000):
        
        # 速度控制器 输入目标速度，输出期望姿态和推力
        v_controller.set_target_velocity_world(np.array([1.0, 1.0, 1.0]))  
        v_controller.update_current_velocity_world(kin.get_state().velocity)
        v_controller.update_current_attitude(kin.get_state().euler)

        # 动力学模型 输入期望姿态和推力，输出加速度和角加速度
        dyn.update_force(v_controller.calculate_force(mass))
        dyn.update_torque(v_controller.calculate_torque(inertia))
        acceleration = dyn.calculate_acceleration(mass, kin.get_state().euler)
        angular_acceleration = dyn.calculate_angular_acceleration(inertia)
        
        # 运动学模型 输入加速度和角加速度，更新位置和姿态
        kin.update_position(acceleration, dt)
        kin.update_attitude_quaternion(angular_acceleration, dt)
        print(f"Position: {kin.get_state().position}, euler: {kin.get_state().euler}")
        # 保存位置
        points = np.vstack((points, kin.get_state().position))
    # 绘制无人机轨迹
    plot_trajectory(points, title='Drone Trajectory')
    
if __name__ == "__main__":
    main()